package com.alibaba.datax.plugin.writer.dmwriter;

import com.alibaba.datax.common.util.Configuration;
import com.alibaba.datax.plugin.rdbms.writer.Constant;
import com.alibaba.druid.pool.DruidDataSource;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.*;

public class DmMetadataUtil {
    private static final Logger LOG = LoggerFactory.getLogger(DmMetadataUtil.class);

    private DruidDataSource dataSource;
    private String schema;
    private Configuration config;
    private String oneTable;

    public DmMetadataUtil(Configuration config) {
        this.config = config;
        String url = config.getString("jdbcUrl");
        if (StringUtils.isBlank(url)) {
            String format = String.format("%s[0].%s", Constant.CONN_MARK, Key.JDBC_URL);
            url = config.getString(format);
        }

        String username = config.getString("username");
        String password = config.getString("password");
        this.schema = getDMSchema(url);
        oneTable = config.getString(String.format("%s[0].%s[0]", Constant.CONN_MARK, Key.TABLE));
        dataSource = new DruidDataSource();
        dataSource.setUrl(url);
        dataSource.setUsername(username);
        dataSource.setPassword(password);
        dataSource.setDriverClassName(Key.DM_JDBC_DRIVER);
        dataSource.setInitialSize(1);
        dataSource.setMaxActive(5);
    }

    /**
     * 获取表的主键列名列表
     */
    public List<String> getPrimaryKeys(String tableName) {
        List<String> primaryKeys = new ArrayList<>();

        try (Connection conn = dataSource.getConnection()) {
            DatabaseMetaData metaData = conn.getMetaData();

            // 达梦数据库表名需要大写
            try (ResultSet rs = metaData.getPrimaryKeys(null, schema, tableName.toUpperCase().replace(schema.toUpperCase() + ".", ""))) {
                while (rs.next()) {
                    String columnName = rs.getString("COLUMN_NAME");
                    short keySeq = rs.getShort("KEY_SEQ");
                    primaryKeys.add(columnName);
                    LOG.debug("表 {} 主键: {}, 序列: {}", tableName, columnName, keySeq);
                }
            }

        } catch (SQLException e) {
            LOG.error("获取表 {} 主键失败", tableName, e);
            throw new RuntimeException("获取主键失败: " + tableName, e);
        }

        return primaryKeys;
    }

    /**
     * 批量获取多个表的主键信息
     */
    public Map<String, List<String>> getBatchPrimaryKeys(List<String> tableNames) {
        Map<String, List<String>> result = new HashMap<>();

        for (String tableName : tableNames) {
            try {
                List<String> primaryKeys = getPrimaryKeys(tableName);
                result.put(tableName, primaryKeys);
            } catch (Exception e) {
                LOG.warn("获取表 {} 主键失败: {}", tableName, e.getMessage());
                result.put(tableName, Collections.emptyList());
            }
        }

        return result;
    }

    /**
     * 获取第一个主键（用于分片键）
     */
    public String getFirstPrimaryKey(String tableName) {
        List<String> primaryKeys = getPrimaryKeys(tableName);
        return primaryKeys.isEmpty() ? null : primaryKeys.get(0);
    }

    public void close() {
        if (dataSource != null) {
            dataSource.close();
        }
    }

    /**
     * 获取达梦模式名 - 当SYSDBA与实际不一致会默认找SYSDBA模式下的表，会找不到
     *
     * @param jdbcUrl 参数 {@link String}
     * @return {@link String}
     * @author xiaohutu 2025-09-22 09:20:45
     */
    public static String getDMSchema(String jdbcUrl) {
        if (StringUtils.isNotBlank(jdbcUrl)) {
            int schemaindex = jdbcUrl.indexOf("schema=");
            if (schemaindex > 0) {
                String substring = jdbcUrl.substring(schemaindex + 1);
                String string = substring.split("&")[0];
                return string.split("=")[1];
            } else {
                int i = jdbcUrl.lastIndexOf("/");
                String substring = jdbcUrl.substring(i + 1);
                return substring.split("\\?")[0];
            }

        }
        return null;
    }
}
